"""
此模块使用 FastAPI 框架将训练好的 GNN 模型封装为 API 服务。
提供单个样本预测的接口，并自动生成 API 文档。
"""
import torch
from torch_geometric.loader import DataLoader
from torch_geometric.data import Data
from fastapi import FastAPI
import os
from model import ZEOEGNN4

app = FastAPI()

# 加载数据集
train_dataset = torch.load('./dataset/train_dataset.pt', weights_only=False)
test_dataset = torch.load('./dataset/test_dataset.pt', weights_only=False)
train_loader = DataLoader(train_dataset, batch_size=1, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)

device = 'cpu'
# 模型参数
n_layers = 1             # GNN层数
node_dim = 7             # 节点特征维度
out_node_nf = 227        # 输出节点特征数/分类类别数
pos_dim = 3              # 位置编码维度
edge_attr_dim = 3        # 边属性维度
hidden_f = 128           # 隐藏层维度
craft_dim = 29           # CRAFT特征维度
craft_hidden = 58        # CRAFT隐藏层维度
craft_out_feat = 29      # CRAFT输出特征数

model = ZEOEGNN4(
    n_layers=n_layers,
    node_dim=node_dim,
    out_node_nf=out_node_nf,
    pos_dim=pos_dim,
    edge_attr_dim=edge_attr_dim,
    hidden_f=hidden_f,
    craft_dim=craft_dim,
    craft_hidden=craft_hidden,
    craft_out_feat=craft_out_feat
)

model_path = './model_epoch_394_best_acc_0.9887.pth'
if os.path.exists(model_path):
    model.load_state_dict(torch.load(model_path, map_location=device))
    model.eval()
    print(f"模型已加载: {model_path}")
else:
    print(f"模型文件 {model_path} 不存在，请检查路径。")


def predict_single(data):
    """
    对单个数据样本进行预测。

    参数:
    data: 单个数据样本，应该包含以下属性:
          - x: 节点特征
          - pos: 位置信息
          - edge_index: 边索引
          - edge_attr: 边特征
          - batch: 批次信息
          - ecfp: ECFP特征
          - maccs: MACCS特征
          - craft_feat: craft特征

    返回:
    dict: 包含预测结果的字典
          - 'prediction': 预测的类别标签
          - 'probabilities': 各类别的概率分布
          - 'confidence': 预测的置信度（最大概率值）
    """
    with torch.no_grad():
        data = data.to(device)
        out = model(
            node_attr=data.x,
            pos=data.pos,
            edge_index=data.edge_index,
            edge_attr=data.edge_attr,
            batch=data.batch,
            ecfp=data.ecfp,
            maccs=data.maccs,
            craft_feat=data.craft_feat
        )
        probabilities = torch.softmax(out, dim=1)
        prediction = torch.argmax(probabilities, dim=1)
        confidence = torch.max(probabilities, dim=1)[0]

        probabilities_np = probabilities.cpu().numpy()
        prediction_np = prediction.cpu().numpy()
        confidence_np = confidence.cpu().numpy()


        # return {
        #     'prediction': prediction_np[0],
        #     'probabilities': probabilities_np[0],
        #     'confidence': confidence_np[0]
        # }
        
        # 转换为 int 类型
        prediction_int = int(prediction_np[0])  # 类别索引，直接转换为 int
        # probabilities_int = int(probabilities_np[0])  # 概率乘以 100 并取整
        confidence_int = int(confidence_np[0])  # 置信度乘以 100 并取整

        return {
            'prediction': prediction_int,
            # 'probabilities': probabilities_int,
            'confidence': confidence_int
        }


# class PredictionRequest(BaseModel):
#     """
#     定义预测请求的数据模型。
#     包含预测所需的所有数据特征，与 Data 对象格式一致。
#     """
#     x: list                 # 节点特征: [n1, 7], float
#     edge_index: list       # 边索引: [2, n2], long
#     edge_attr: list        # 边属性: [n2, 3], float
#     pos: list              # 节点位置: [n1, 3], float
#     z: list                # 节点属性: [n1], long
#     smiles: str            # SMILES 字符串
#     maccs: list            # MACCS 指纹: [1, 167], long
#     ecfp: list             # ECFP 指纹: [1, 2048], long
#     y: int                 # 标签: int
#     craft_feat: list       # CRAFT 特征: [1, 29], float


@app.post("/predict")
def predict(request_data:dict):
    """
    FastAPI API 端点，用于处理预测请求。
    接收包含数据样本的请求，返回预测结果。
    """
    # 构建 Data 对象
    data_obj = Data(
        x=torch.tensor(request_data["x"], dtype=torch.float32),
        edge_index=torch.tensor(request_data["edge_index"], dtype=torch.int64),
        edge_attr=torch.tensor(request_data["edge_attr"], dtype=torch.float32),
        pos=torch.tensor(request_data["pos"], dtype=torch.float32),
        z=torch.tensor(request_data["z"], dtype=torch.int64),
        smiles=request_data["smiles"],
        maccs=torch.tensor(request_data["maccs"], dtype=torch.float32),
        ecfp=torch.tensor(request_data["ecfp"], dtype=torch.float32),
        y=request_data["y"],
        craft_feat=torch.tensor(request_data["craft_feat"], dtype=torch.float32)
    )
    result = predict_single(data_obj)
    return result


